"""
本工具用于处理视频中的二值化图像，允许用户通过滑动条调整矩形框的大小和位置，并计算该区域内黑色像素的占比。
主要处理任务：面对特殊形状线的识别（如箭头，横线等等）

"""

import cv2
import numpy as np

# 打开视频
cap = cv2.VideoCapture('test.mp4')

# 设置视频大小
cap.set(cv2.CAP_PROP_FRAME_WIDTH, 640)
cap.set(cv2.CAP_PROP_FRAME_HEIGHT, 480)

# 创建一个窗口
cv2.namedWindow('Video')
cv2.namedWindow('Binary Image')

# 初始化滑动条（例如，调整矩形框的x, y坐标和宽度、高度）
cv2.createTrackbar('X', 'Video', 60, 640, lambda x: None)
cv2.createTrackbar('Y', 'Video', 20, 480, lambda x: None)
cv2.createTrackbar('Width', 'Video', 580-60, 640, lambda x: None)
cv2.createTrackbar('Height', 'Video', 440-400, 480, lambda x: None)

# 初始化标定的矩形框
start_point = None
end_point = None
drawing = False


def mouse_callback(event, x, y, flags, param):
    global start_point, end_point, drawing

    if event == cv2.EVENT_LBUTTONDOWN:
        start_point = (x, y)
        drawing = True
    elif event == cv2.EVENT_MOUSEMOVE:
        if drawing:
            end_point = (x, y)
    elif event == cv2.EVENT_LBUTTONUP:
        drawing = False
        end_point = (x, y)


# 设置鼠标回调
cv2.setMouseCallback('Video', mouse_callback)

# 帧计数器
frame_counter = 0

while True:
    # 跳到指定的帧
    cap.set(cv2.CAP_PROP_POS_FRAMES, frame_counter)

    ret, frame = cap.read()
    if not ret:
        break

    frame = cv2.resize(frame, (640, 480))  # 调整图像大小
    # hsv_image = cv2.cvtColor(frame, cv2.COLOR_BGR2HSV)
    # h, s, v = cv2.split(hsv_image)
    # mask = h > 80
    # hsv_image[mask] = (0, 0, 255)
    # frame = cv2.cvtColor(hsv_image, cv2.COLOR_HSV2BGR)

    # 获取滑动条的值
    x = cv2.getTrackbarPos('X', 'Video')
    y = cv2.getTrackbarPos('Y', 'Video')
    width = cv2.getTrackbarPos('Width', 'Video')
    height = cv2.getTrackbarPos('Height', 'Video')

    # 如果存在鼠标框定，使用鼠标选定的框
    if start_point and end_point:
        x, y = start_point
        width, height = end_point[0] - start_point[0], end_point[1] - start_point[1]

    # 绘制矩形框
    cv2.rectangle(frame, (x, y), (x + width, y + height), (0, 255, 0), 2)

    # 转换为HSV颜色空间
    hsv = cv2.cvtColor(frame, cv2.COLOR_BGR2HSV)
    
    # 创建红色掩码（两部分）
    lower_red1 = np.array([0, 100, 100])
    upper_red1 = np.array([15, 255, 255])
    lower_red2 = np.array([165, 100, 100])
    upper_red2 = np.array([179, 255, 255])
    
    mask1 = cv2.inRange(hsv, lower_red1, upper_red1)
    mask2 = cv2.inRange(hsv, lower_red2, upper_red2)
    
    # 合并两个掩码
    mask = cv2.bitwise_or(mask1, mask2)
    
    # 对原始图像和掩码进行位运算
    result = cv2.bitwise_and(frame, frame, mask=mask)
    binary_mask = np.zeros((result.shape[0], result.shape[1]), dtype=np.uint8)
    non_black_pixels = np.any(result > 0, axis=-1)
    binary_mask[non_black_pixels] = 255
    # # 转为灰度图
    # kernel = cv2.getStructuringElement(cv2.MORPH_CROSS, (5, 5))
    # gray = cv2.cvtColor(result, cv2.COLOR_BGR2GRAY)
    # gray = cv2.erode(gray, kernel)  # 腐蚀结果
    # gray = cv2.dilate(gray, kernel)  # 膨胀结果

    # # 设置二值化阈值
    # threshold_value = 30
    # _, binary = cv2.threshold(gray, threshold_value, 255, cv2.THRESH_BINARY)  # 黑色为255，白色为0

    # 提取矩形区域
    roi = binary_mask[y:y + height, x:x + width]

    # 计算黑色像素占比
    black_pixels = cv2.countNonZero(roi)  # 计算黑色像素数
    total_pixels = roi.size
    if total_pixels == 0:
        total_pixels = 1
    black_ratio = black_pixels / total_pixels

    print(f'黑色像素占比: {black_ratio:.2f}, {black_pixels},{total_pixels}')

    # 显示原始视频图像
    cv2.imshow('Video', frame)

    # 显示二值化图像
    cv2.imshow('Binary Image', binary_mask)

    key = cv2.waitKey(1) & 0xFF
    if key == ord('q'):  # 按'q'退出
        break
    elif key == 32:  # 按空格键切换到下一帧
        frame_counter += 1
        if frame_counter >= int(cap.get(cv2.CAP_PROP_FRAME_COUNT)):  # 如果已经是最后一帧，重置
            frame_counter = 0

cap.release()
cv2.destroyAllWindows()


